﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.Xna.Framework.Graphics;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Graphics.PackedVector;
using Microsoft.Xna.Framework.Content;
using System.IO;

namespace FrogEngine
{
	public class StatefulParticleSystem : GameEntityBase3D, IRenderable, IHasContent, IDisposable
	{
		public StatefulParticleSystem(IEngine engine)
			: base(engine)
		{
			this.graphicsDevice = engine.GraphicsDevice;
			this.spriteBatch = engine.SpriteBatch;
			this.effect = new StatefulParticleEffect(engine);

			this.positionTex = new RenderTarget2D(engine.GraphicsDevice, w, h, false, SurfaceFormat.Vector4, DepthFormat.None, 1, RenderTargetUsage.DiscardContents);
			this.positionRenderTarget = new RenderTarget2D(engine.GraphicsDevice, w, h, false, SurfaceFormat.Vector4, DepthFormat.None, 1, RenderTargetUsage.DiscardContents);
			this.tempPositionBuffer = new Vector4[w * h];
			this.velocityTex = new RenderTarget2D(engine.GraphicsDevice, w, h, false, SurfaceFormat.Vector4, DepthFormat.None, 1, RenderTargetUsage.DiscardContents);
			this.velocityRenderTarget = new RenderTarget2D(engine.GraphicsDevice, w, h, false, SurfaceFormat.Vector4, DepthFormat.None, 1, RenderTargetUsage.DiscardContents);
			this.tempVelocityBuffer = new Vector4[w * h];

			this.vertBuff = new StatefulParticleVertex[numParticles * 4];
			this.indBuff = new int[numParticles * 6];
			this.vertexBuffer = new DynamicVertexBuffer(engine.GraphicsDevice, typeof(StatefulParticleVertex), numParticles * 4, BufferUsage.None);
			this.indexBuffer = new IndexBuffer(engine.GraphicsDevice, IndexElementSize.ThirtyTwoBits, numParticles * 6, BufferUsage.None);
			this.InitializeIndices();

			this.cameraService = engine.GetService<ICameraService>();

			Random random = new Random();
			for (int i = 0; i < w; i++)
			{
				for (int j = 0; j < h; j++)
				{
					float x = random.Next(1, 2000);
					float y = random.Next(1, 2000);
					float z = random.Next(1, 2000);
					this.AddParticle(new Vector3(x, y, z), Vector3.Zero, j, i);
				}
			}

			this.positionTex.SetData<Vector4>(tempPositionBuffer);
			this.vertexBuffer.SetData(vertBuff);

			Services.Add(engine.GetService<RenderService>());
			Services.Add(engine.GetService<GameContentService>());

			foreach (var service in Services)
				service.Subscribe(this);
		}

		private const int numParticles = 1000000;
		private const int w = 1000;
		private const int h = 1000;

		private StatefulParticleEffect effect;
		private Texture2D particleTex;
		private RenderTarget2D positionTex;
		private Vector4[] tempPositionBuffer;
		private RenderTarget2D positionRenderTarget;
		private RenderTarget2D velocityTex;
		private Vector4[] tempVelocityBuffer;
		private RenderTarget2D velocityRenderTarget;

		private StatefulParticleVertex[] vertBuff;
		private DynamicVertexBuffer vertexBuffer;

		private int[] indBuff;
		private IndexBuffer indexBuffer;

		private GraphicsDevice graphicsDevice;
		private SpriteBatch spriteBatch;
		private ICameraService cameraService;

		public void Dispose()
		{
			positionTex.Dispose();
			positionRenderTarget.Dispose();
			velocityTex.Dispose();
			velocityRenderTarget.Dispose();
			vertexBuffer.Dispose();
			indexBuffer.Dispose();
			indBuff = null;
			vertBuff = null;
			tempPositionBuffer = null;
			tempVelocityBuffer = null;
			GC.Collect();
		}

		public void Update(GameTime gameTime)
		{
		}

		private int firstFree = 0;
		public void AddParticle(Vector3 position, Vector3 velocity, int x, int y)
		{
			float index = firstFree / 4;
			vertBuff[firstFree + 0] = new StatefulParticleVertex { UV = new Short2(x, y), Corner = new Short2(-1, -1) };
			vertBuff[firstFree + 1] = new StatefulParticleVertex { UV = new Short2(x, y), Corner = new Short2(1, -1) };
			vertBuff[firstFree + 2] = new StatefulParticleVertex { UV = new Short2(x, y), Corner = new Short2(1, 1) };
			vertBuff[firstFree + 3] = new StatefulParticleVertex { UV = new Short2(x, y), Corner = new Short2(-1, 1) };
			tempPositionBuffer[(int)index] = new Vector4(position, 0);
			firstFree += 4;
			if (firstFree > numParticles * 4)
				firstFree = 0;
		}

		public void Render(GameTime gameTime)
		{
			var renderService = this.GetService<RenderService>();
			var camera = this.cameraService.GetActiveCamera();

			StatefulParticleShaderState state = new StatefulParticleShaderState();
			state.World = Matrix.Identity;
			state.View = camera.View;
			state.Projection = camera.Projection;
			state.VelocityTex = this.velocityTex;
			state.PositionTex = this.positionTex;
			state.ParticleTex = this.particleTex;
			state.dt = (float)gameTime.ElapsedGameTime.TotalSeconds;
			state.ViewPortScale = new Vector2(0.5f / graphicsDevice.Viewport.AspectRatio, -0.5f);
			state.globalForce = new Vector3(0, -100, 0);
			state.size = 2;
			state.attractiveLocalCount = 2;
			state.attractiveLocalMagnitudes = new float[4] { 600000f, 600000f, 1f, 1f };
			state.attractiveLocals = new Vector3[4] { camera.Position + camera.Forward * 2000, new Vector3(1000, 1000, 1000), Vector3.Zero, Vector3.Zero };
			state.repulsiveLocalCount = 0;
			state.repulsiveLocalMagnitudes = new float[4] { 400000f, 1f, 1f, 1f };
			state.repulsiveLocals = new Vector3[4] { camera.Position, Vector3.Zero, Vector3.Zero, Vector3.Zero };
			this.effect.SetShaderParameters(state);

			Viewport viewport = this.graphicsDevice.Viewport;
			Matrix projection = Matrix.CreateOrthographicOffCenter(0, w, h, 0, 0, 1);
			Matrix halfPixelOffset = Matrix.CreateTranslation(-0.5f, -0.5f, 0);
			Matrix combined = halfPixelOffset * projection;

			this.effect.ActualEffect.Parameters["MatrixTransform"].SetValue(combined);
			
			// calculate new velocities with shader
			effect.StartVelocityTechnique();
			graphicsDevice.SetRenderTarget(velocityRenderTarget);
			spriteBatch.Begin(SpriteSortMode.Immediate, BlendState.Opaque, null, null, null, this.effect.ActualEffect);
			spriteBatch.Draw(positionTex, new Rectangle(0, 0, w, h), Color.White);
			spriteBatch.End();
			graphicsDevice.SetRenderTarget(positionRenderTarget);

			this.effect.ActualEffect.Parameters["VelocityTex"].SetValue(velocityRenderTarget);
			this.effect.ActualEffect.Parameters["PositionTex"].SetValue(positionTex);

			// calculate new positions with shader
			effect.StartPositionTechnique();
			spriteBatch.Begin(SpriteSortMode.Immediate, BlendState.Opaque, null, null, null, this.effect.ActualEffect);
			spriteBatch.Draw(positionTex, new Rectangle(0, 0, w, h), Color.White);
			spriteBatch.End();
			graphicsDevice.SetRenderTarget(renderService.SceneTarget);

			this.effect.ActualEffect.Parameters["PositionTex"].SetValue(positionRenderTarget);

			// Restore the vertex buffer contents if the graphics device was lost.
			if (vertexBuffer.IsContentLost)
			{
				vertexBuffer.SetData(vertBuff);
			}

			this.graphicsDevice.SetVertexBuffer(vertexBuffer);
			this.graphicsDevice.Indices = indexBuffer;
			this.graphicsDevice.DepthStencilState = DepthStencilState.DepthRead;
			this.graphicsDevice.BlendState = BlendState.AlphaBlend;
			this.graphicsDevice.RasterizerState = RasterizerState.CullNone;
			this.effect.StartRenderTechnique();
			foreach (var pass in this.effect.CurrentTechnique.Passes)
			{
				pass.Apply();
				this.graphicsDevice.DrawIndexedPrimitives(PrimitiveType.TriangleList, 0, 0, numParticles * 2, 0, numParticles);
				this.graphicsDevice.DrawIndexedPrimitives(PrimitiveType.TriangleList, 0, numParticles * 2, numParticles * 2, numParticles * 3, numParticles);
			}

			// swap velocity rendertargets
			RenderTarget2D tempVel = velocityRenderTarget;
			velocityRenderTarget = velocityTex;
			velocityTex = tempVel;

			// swap position rendertargets
			RenderTarget2D tempPos = positionRenderTarget;
			positionRenderTarget = positionTex;
			positionTex = tempPos;
		}

		#region [ Private Methods ]

		private void InitializeIndices()
		{
			int partStride = 4; // 4 vertices per particle
			for (int j = 0; j < numParticles; j++)
			{
				int k = j * 6;
				int i = partStride * j;
				indBuff[k + 0] = i + 0;
				indBuff[k + 1] = i + 1;
				indBuff[k + 2] = i + 2;
				indBuff[k + 3] = i + 0;
				indBuff[k + 4] = i + 2;
				indBuff[k + 5] = i + 3;
			}

			indexBuffer.SetData(indBuff);
		}

		#endregion

		#region [ Content Management ]

		private ContentManager contentManager;
		public Microsoft.Xna.Framework.Content.ContentManager ContentManager
		{
			get { return contentManager; }
			set { contentManager = value; }
		}

		public IEnumerable<string> AssetNames()
		{
			return new List<string> { "GenericParticle" };
		}

		public void LoadContent()
		{
			this.particleTex = this.ContentManager.Load<Texture2D>("GenericParticle");
		}

		#endregion
	}
}
